#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Kailash Budhathoki (kailash.buki@gmail.com)
"""Synthetic class labelled dataset generator which randomly injects pattern too. The gen_rand.json config file must be configured before running this file.
    "injection_type": [mixed | local | common | nopat]
"""

import json
import os
import random
import sys


def rand_cardinality():
    # 4 to 8 singletons in a pattern
    return random.randint(4, 8)

def rand_num_pat():
    # 4 to 10 patterns in random
    return random.randint(4, 10)

def rand_freq():
    # 10% to 30% frequency
    return random.choice([0.1, 0.2, 0.3])

def get_cardinality(itemset):
    return len(itemset)


class RandomDataGenerator(object):
    def __init__(self, config_file):
        self.dataset = ''
        self.inject = dict()
        self.stats = dict()
        self.common_patterns = []
        self.local_patterns = []
        self.injected_alphabets = set()
        self.remaining_alphabets = set()
        self.read_config(config_file)
        self.generate_alphabets()
        self.generate_class_labels()

    def read_config(self, config_file):
        with open(config_file) as conf_file:
            self.config = json.load(conf_file)

    def generate_alphabets(self):
        self.alphabets = [alphabet for alphabet in xrange(self.config['num_alphabets'])]

    def generate_class_labels(self):
        self.class_labels = []
        max_alphabet = max(self.alphabets)
        next_label = max_alphabet + 1
        for i in xrange(self.config['num_splits']):
            self.class_labels.append(next_label)
            self.inject[i] = dict()
            self.stats[i] = dict()
            next_label += 1

    def copy_stats(self):
        for i in xrange(self.config['num_splits']):
            for pattern, desired_freq in self.inject[i].iteritems():
                planted_count = 0
                support = 0
                self.stats[i][pattern] = [desired_freq, planted_count, support]

    def update_remaining_alphabets(self):
        self.remaining_alphabets = set(self.alphabets).difference(self.injected_alphabets)

    def create_output_dir(self, injection_type):
        self.output_dir = "synth_%s" % injection_type
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def clear_stats(self):
        for i in xrange(self.config['num_splits']):
            self.inject[i] = dict()
            self.stats[i] = dict()

    def clear_variables(self):
        self.clear_stats()
        self.dataset = ''
        self.common_patterns = []
        self.local_patterns = []
        self.injected_alphabets = set()
        self.remaining_alphabets = set()

    def generate(self):
        for injection_type in self.config['injection_type']:
            self.get_injection_func(injection_type)()
            self.update_remaining_alphabets()
            self.copy_stats()
            self.create_dataset()
            self.create_output_dir(injection_type)
            self.ouput_dataset(injection_type)
            self.output_stats(injection_type)
            self.clear_variables()

    def create_split(self, split_id):
        split = ''
        class_label = self.class_labels[split_id]
        split_inject = self.inject[split_id]
        split_stats = self.stats[split_id]
        num_rows = self.config['num_rows_per_split'][split_id]
        rows_generated = 0
        while rows_generated < num_rows:
            row = set()

            for itemset, freq in split_inject.iteritems():
                if random.random() < freq:
                    row = row.union(itemset)
                    split_stats[itemset][-2] += 1       # planted count

            for alphabet in self.remaining_alphabets:
                if alphabet not in row and random.random() < self.config['uniform_noise']:
                    row.add(alphabet)

            for itemset, freq in split_inject.iteritems():
                if set(itemset) == row:
                    split_stats[itemset][-1] += 1       # support, itemset is a transaction, we didn't add any noise

            if len(row) == 0:
                continue

            row = str(class_label) + ' ' + ' '.join(str(alphabet) for alphabet in list(sorted(row))) + '\n'
            split += row
            rows_generated += 1

        return split

    def create_dataset(self):
        for i in xrange(self.config['num_splits']):
            split = self.create_split(i)
            self.dataset += split

    def ouput_dataset(self, injection_type):
        dataset_file = os.path.join(self.output_dir, "synth_%s.dat" % injection_type)
        with open(dataset_file, 'w') as op:
            op.write(self.dataset)

    def output_stats(self, injection_type):
        console = sys.stdout
        stats_file = os.path.join(self.output_dir, "synth_%s.stats" % injection_type)
        sys.stdout = open(stats_file, 'w')
        self.pretty_print_injected_patterns()
        sys.stdout = console

    def get_injection_func(self, injection_type):
        if injection_type == 'local':
            return self.generate_injection_pattern_local
        elif injection_type == 'common':
            return self.generate_injection_pattern_common
        elif injection_type == 'mixed':
            return self.generate_injection_pattern_mixed
        elif injection_type == 'nopat':
            return self.generate_injection_pattern_nopat
        else:
            print 'Unknown injection type!!!'
            exit(0)

    def generate_injection_pattern_local(self):
        for i in xrange(self.config['num_splits']):
            num_local_pat = rand_num_pat()
            pat_generated = 0
            while pat_generated < num_local_pat:
                cardinality = rand_cardinality()
                freq = rand_freq()
                itemset = set()
                while get_cardinality(itemset) < cardinality:
                    alphabet = random.choice(self.alphabets)
                    if alphabet not in itemset and alphabet not in self.injected_alphabets:
                        self.injected_alphabets.add(alphabet)
                        itemset.add(alphabet)

                if itemset not in self.local_patterns:
                    self.inject[i][tuple(itemset)] = freq
                    self.local_patterns.append(tuple(itemset))
                    pat_generated += 1


    def generate_injection_pattern_common(self):
        num_common_pat = rand_num_pat()
        pat_generated = 0
        while pat_generated < num_common_pat:
            cardinality = rand_cardinality()
            freq = rand_freq()
            itemset = set()
            while get_cardinality(itemset) < cardinality:
                alphabet = random.choice(self.alphabets)
                if alphabet not in self.injected_alphabets and alphabet not in self.injected_alphabets:
                    self.injected_alphabets.add(alphabet)
                    itemset.add(alphabet)

            if itemset not in self.common_patterns:
                self.common_patterns.append(itemset)
                pat_generated += 1
                for i in xrange(self.config['num_splits']):
                    self.inject[i][tuple(itemset)] = freq


    def generate_injection_pattern_mixed(self):
        self.generate_injection_pattern_common()
        self.generate_injection_pattern_local()

    def generate_injection_pattern_nopat(self):
        pass

    def pretty_print_injected_patterns(self):
        col1_width = 40
        col2_width = 7
        col3_width = 7
        col4_width = 7
        col5_width = 7
        for i in xrange(self.config['num_splits']):
            repeat = 105
            print
            print "-" * repeat
            print "\t\t\t\t\t\tSPLIT ", i+1
            print "-" * repeat
            print "%s  \t|%s\t|%s\t|%s\t|%s" % ("Planted ItemSet".rjust(col1_width), "Desired Freq".rjust(col2_width), "Planted Freq".rjust(col3_width), "Planted Count".rjust(col4_width), "Support".rjust(col5_width))
            print "-" * repeat
            for pattern, stats in self.stats[i].iteritems():
                pattern = sorted(pattern)
                desired_freq, planted_count, support = stats
                num_rows = self.config['num_rows_per_split'][i]
                planted_freq = round(planted_count / float(num_rows), 2)
                if set(pattern) in self.common_patterns:
                    print '%s *\t|%s\t|%s\t|%s\t|%s' %(str(pattern).rjust(col1_width), str(desired_freq).rjust(col2_width), str(planted_freq).rjust(col3_width), str(planted_count).rjust(col4_width), str(support).rjust(col5_width))
                else:
                    print '%s  \t|%s\t|%s\t|%s\t|%s' %(str(pattern).rjust(col1_width), str(desired_freq).rjust(col2_width), str(planted_freq).rjust(col3_width), str(planted_count).rjust(col4_width), str(support).rjust(col5_width))
            print "-" * repeat

if __name__ == '__main__':
    rg = RandomDataGenerator('gen_rand.json')
    rg.generate()
